use crate::cjson::cjson_print_from_value;
use crate::verifier::ItrusteePlugin;
use base64::engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD};
use base64::Engine;
use openssl::{
    asn1::Asn1Time,
    hash::MessageDigest,
    pkey::{PKey, Public},
    rsa::Padding,
    sign::{RsaPssSaltlen, Verifier},
    x509::X509,
};
use plugin_manager::PluginError;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::cmp::Ordering::{Greater, Less};
use std::collections::HashMap;

/// ITrusteeEvidence Root Structure
#[derive(Debug, Clone)]
pub struct ITrusteeEvidence {
    pub json_value: Value,
    pub report_sign: ReportSign,
    pub ak_cert: AkCert,
    pub payload: Payload,
    pub handler: String,
}

/// Payload of ITrusteeEvidence
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Payload {
    pub version: String,
    pub timestamp: String,
    pub nonce: String,
    pub scenario: String,
    pub uuid: String,
    pub requester_uuid: String,
    pub sign_alg: String,
    pub hash_alg: String,
    pub ta_mem: String,
    pub ta_img: String,
    pub ta_attr: String,
    pub tcb: String,

    /// Reserved for unknown fields
    #[serde(flatten)]
    pub extra: HashMap<String, Value>,
}

/// ReportSign Enum
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum ReportSign {
    SceNoAs { sce_no_as: String },        // base64
}

/// AkCert Enum
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum AkCert {
    SceNoAs { sce_no_as: AkCertNoAs },
}

/// AK certificate structure for NO AS scenario: contains payload and signature
#[derive(Debug, Clone, Deserialize)]
pub struct AkCertNoAs {
    pub payload: AkCertNoAsPayload,
    pub signature: AkCertNoAsSignature,
}

/// payload: at least contains ak_pub, qta_img, qta_mem
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AkCertNoAsPayload {
    pub ak_pub: AkPub,

    /// qta image and memory hashes, both are base64url
    pub qta_img: String,
    pub qta_mem: String,

    /// Reserved fields
    #[serde(flatten)]
    pub extra: HashMap<String, Value>,
}

/// DRK signature and certificate: both are base64url
#[derive(Debug, Clone, Deserialize)]
pub struct AkCertNoAsSignature {
    pub drk_sign: String,
    pub drk_cert: String,

    /// Reserved fields
    #[serde(flatten)]
    pub extra: HashMap<String, Value>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AkPub {
    #[serde(rename = "kty")]
    pub key_type: String,

    #[serde(flatten)]
    pub extra: HashMap<String, Value>,
}

#[derive(Debug, Deserialize)]
pub struct Logs {
    pub log_type: String,
    pub log_data: String,
}

impl ITrusteeEvidence {
    /// JSON format for evidence input
    /// {
    ///     "log": {
    ///         "log_type": "ItrusteeLog",
    ///         "log_data": "<base64 of inner JSON>"
    ///     }
    /// }
    pub fn from_json_value(value: &Value) -> Result<Self, PluginError> {
        // 1. Find evidence.logs[].log_data
        let logs = value
            .get("logs")
            .and_then(|v| v.as_array())
            .ok_or_else(|| PluginError::InputError("evidence.logs not exist".to_string()))?;
        
        let mut b64_log_data: Option<&str> = None;
        for log in logs {
            if log.get("log_type").and_then(|v| v.as_str()) == Some("ItrusteeLog") {
                b64_log_data = log.get("log_data").and_then(|v| v.as_str());
                break;
            }
        }
        let b64_log_data = b64_log_data.ok_or_else(|| PluginError::InputError("missing log which log_type is 'ItrusteeLog'".to_string()))?;

        // 2. base64 decode JSON
        let raw = STANDARD.decode(b64_log_data).map_err(|e| PluginError::InputError(format!("base64 decode failed: {e}")))?;
        let inner: Value =
            serde_json::from_slice(&raw).map_err(|e| PluginError::InputError(format!("json decode failed: {e}")))?;

        // 3. payload / handle
        let handler_json =
            inner.get("handler").cloned().ok_or_else(|| PluginError::InputError("missing handler".to_string()))?;
        let payload_json = Self::get_payload_json(&inner).cloned()?;

        let payload: Payload =
            serde_json::from_value(payload_json).map_err(|e| PluginError::InputError(format!("payload parse: {e}")))?;
        Self::check_scenario(payload.scenario.as_str())?;
        let handler: String = serde_json::from_value(handler_json)
            .map_err(|e| PluginError::InputError(format!("parse handler failed: {e}")))?;

        // 4. report_sign / akcert
        let report_sign_json = Self::get_report_sign_json(&inner).cloned()?;
        let akcert_json = Self::get_ak_cert_json(&inner).cloned()?;

        let report_sign = serde_json::from_value(report_sign_json)
            .map_err(|e| PluginError::InputError(format!("parse report_sign(no_as) failed: {e}")))?;
        let ak_cert: AkCert = serde_json::from_value(akcert_json)
            .map_err(|e| PluginError::InputError(format!("parse ak cert(as_no_daa) failed: {e}")))?;

        let json_value = inner.clone();

        Ok(ITrusteeEvidence { json_value, report_sign, ak_cert, payload, handler })
    }

    pub async fn verify(
        &mut self,
        _user_id: &str,
        _node_id: Option<&str>,
        nonce: Option<&[u8]>,
        _plugin: &ItrusteePlugin,
    ) -> Result<Value, PluginError> {
        // 1.nonce verification
        self.verify_nonce(nonce)?;

        // 2.signature verification: verify according to different scenarios
        self.verify_signature()
    }

    fn verify_nonce(&mut self, nonce: Option<&[u8]>) -> Result<(), PluginError> {
        if let Some(expected_nonce) = nonce {
            // Get nonce from payload and decode it
            let report_nonce = URL_SAFE_NO_PAD
                .decode(&self.payload.nonce)
                .map_err(|e| PluginError::InputError(format!("decode nonce failed: {e}")))?;

            // Compare nonce length and content
            if report_nonce.len() != expected_nonce.len() || report_nonce != expected_nonce {
                return Err(PluginError::InputError("nonce verification failed".to_string()));
            }
        }
        Ok(())
    }

    fn verify_signature(&mut self) -> Result<Value, PluginError> {
        match self.payload.scenario.as_str() {
            "sce_no_as" => self.verify_no_as(),
            other => Err(PluginError::InputError(format!("unsupported scenario: {other}"))),
        }
    }
}

/// Common methods used by three scenarios
impl ITrusteeEvidence {
    pub fn check_scenario(sce: &str) -> Result<(), PluginError> {
        match sce {
            "sce_no_as" => Ok(()),
            other => Err(PluginError::InputError(format!("unsupported scenario: {other}"))),
        }
    }
    pub(crate) fn get_payload_json(qca_report: &Value) -> Result<&Value, PluginError> {
        qca_report.get("payload").ok_or_else(|| PluginError::InputError("missing field 'payload'".to_string()))
    }

    pub(crate) fn get_report_sign_json(qca_report: &Value) -> Result<&Value, PluginError> {
        qca_report.get("report_sign").ok_or_else(|| PluginError::InputError("missing field 'report_sign'".to_string()))
    }

    pub(crate) fn get_ak_cert_json(qca_report: &Value) -> Result<&Value, PluginError> {
        qca_report.get("akcert").ok_or_else(|| PluginError::InputError("missing field 'akcert'".to_string()))
    }

    pub(crate) fn get_ak_payload_json(qca_report: &Value) -> Result<&Value, PluginError> {
        let sce_no_as = Self::get_ak_cert_json(qca_report)?.get("sce_no_as").ok_or_else(|| {
            PluginError::InputError("missing field 'sce_no_as' while extracting ak payload".to_string())
        })?;
        sce_no_as.get("payload").ok_or_else(|| PluginError::InputError("missing field 'payload'".to_string()))
    }

    /// Verify certificate validity
    pub(crate) fn verify_cert_validity(cert: &X509, cert_type: &str) -> Result<(), PluginError> {
        let now = Asn1Time::days_from_now(0)
            .map_err(|e| PluginError::InternalError(format!("Failed to get current time: {}", e)))?;

        if now.compare(cert.not_before()).map_err(|e| PluginError::InternalError(e.to_string()))? != Greater ||
           now.compare(cert.not_after()).map_err(|e| PluginError::InternalError(e.to_string()))? != Less {
            return Err(PluginError::InternalError(format!("{} certificate is expired or not yet valid", cert_type)));
        }
        Ok(())
    }

    /// Encode message payload with b64 as data to be signed
    pub(crate) fn encode_payload(&self) -> Result<String, PluginError> {
        let report_payload_json = Self::get_payload_json(&self.json_value)?;
        let sign_data = URL_SAFE_NO_PAD.encode(cjson_print_from_value(report_payload_json)?);
        Ok(sign_data)
    }

    pub(crate) fn verify_payload(&self, report_sign: &String, pubkey: PKey<Public>) -> Result<(), PluginError> {
        let sign_data = self.encode_payload()?;
        Self::verify_sign_by_key(&sign_data, report_sign, pubkey)?;
        Ok(())
    }

    pub(crate) fn verify_cert_by_cert(issuer_cert: &X509, subject_cert: &X509) -> Result<(), PluginError> {
        // Extract parent certificate public key
        let ca_pubkey = issuer_cert
            .public_key()
            .map_err(|e| PluginError::InternalError(format!("load pub key from root CA failed: {e}")))?;

        // Use public key to verify certificate
        match subject_cert.verify(&ca_pubkey) {
            Ok(false) | Err(_) => return Err(PluginError::InternalError("cert verify failed".to_string())),
            _ => {},
        }
        Ok(())
    }

    /// Verify signature by public key
    pub(crate) fn verify_sign_by_key(
        sign_data_b64url: &String,
        ak_sign_b64url: &String,
        pub_key: PKey<Public>,
    ) -> Result<(), PluginError> {
        // 1) Base64URL decode "signature value"
        let sig = URL_SAFE_NO_PAD
            .decode(ak_sign_b64url)
            .map_err(|e| PluginError::InternalError(format!("decode signature failed: {e}")))?;

        // 2) Construct PSS verifier: SHA256 + MGF1(SHA256) + saltlen=AUTO (equivalent to C's sLen=-2)
        let mut verifier = Verifier::new(MessageDigest::sha256(), &pub_key)
            .map_err(|e| PluginError::InternalError(format!("create verifier failed: {e}")))?;
        verifier
            .set_rsa_padding(Padding::PKCS1_PSS)
            .map_err(|e| PluginError::InternalError(format!("RSA padding failed: {e}")))?;
        verifier
            .set_rsa_mgf1_md(MessageDigest::sha256())
            .map_err(|e| PluginError::InternalError(format!("set mgf1 md failed: {e}")))?;
        verifier
            .set_rsa_pss_saltlen(RsaPssSaltlen::custom(-2))
            .map_err(|e| PluginError::InternalError(format!("set saltlen failed: {e}")))?;

        // 3) Update verifier with signature data
        verifier
            .update(sign_data_b64url.as_bytes())
            .map_err(|e| PluginError::InternalError(format!("update sign data failed: {e}")))?;

        match verifier.verify(&sig) {
            Ok(true) => Ok(()),
            _ => Err(PluginError::InternalError("sign verify failed".to_string())),
        }
    }
}
